#include "register/op_def_registry.h"
#include "graph/utils/type_utils.h"
#include "tiling/platform/platform_ascendc.h"
#include <cmath>
#include <iostream>

#include "nll_loss_tiling.h"

namespace optiling {
static ge::graphStatus TilingFunc(gert::TilingContext* context)
{

    NLLLossTilingData tiling; 
    constexpr bool DOUBLE_BUFFER_ENABLE = true;
    platform_ascendc::PlatformAscendC ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
    uint32_t CORE_NUM;

    const auto reduce_mode = context->GetAttrs()->GetStr(0);
    const int ignore_index = *context->GetAttrs()->GetInt(1);
    const auto& x_shape = context->GetInputShape(0)->GetStorageShape();
    uint32_t N, C, O;
    if (x_shape.GetDimNum()==1)
    {
        N = 1;
        C = x_shape.GetDim(0);
        O = 1;
    }else if(x_shape.GetDimNum()==2){
        N = x_shape.GetDim(0);
        C = x_shape.GetDim(1);
        O = 1;
    }else{
        N = x_shape.GetDim(0);
        C = x_shape.GetDim(1);
        O = x_shape.GetShapeSize()/N/C;
    }
    
    if(reduce_mode[0] == 's'){
        context->SetTilingKey(0);
    }else if(reduce_mode[0] == 'm'){
        context->SetTilingKey(1);
    }
    
    uint32_t fn,fb,tb;
    CORE_NUM = ascendcPlatform.GetCoreNumAiv();
    if(reduce_mode[0] == 'm'){
        CORE_NUM = 1;
    }else{
        if(N==1){
            fn = 1;
            fb = 1;
            tb = 0;
            CORE_NUM = 1;
        }else if(N<512){
            fn = 9;
            fb = 1;
            tb = 0;
            CORE_NUM = 9;
        }else{
            fn = 39;
            fb = 26;
            tb = 10;
            CORE_NUM = 40;
        }
    }
    tiling.set_N(N);                           
    tiling.set_C(C); 
    tiling.set_fn(fn); 
    tiling.set_fb(fb); 
    tiling.set_tb(tb); 
    tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
    context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());
    size_t *currentWorkspace = context->GetWorkspaceSizes(1);
    currentWorkspace[0] = 0;
    
    context->SetBlockDim(CORE_NUM);

    return ge::GRAPH_SUCCESS;
}
}

namespace ge {
static ge::graphStatus InferShape(gert::InferShapeContext* context)
{
    const gert::Shape* x1_shape = context->GetInputShape(0);
    gert::Shape* y_shape = context->GetOutputShape(0);
    *y_shape = *x1_shape;
    return GRAPH_SUCCESS;
}
}


namespace ops {
class NLLLoss : public OpDef {
public:
    explicit NLLLoss(const char* name) : OpDef(name)
    {
        this->Input("x")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT})
            .Format({ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND});
        this->Input("target")
            .ParamType(REQUIRED)
            .DataType({ge::DT_INT32})
            .Format({ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND});
        this->Input("weight")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT})
            .Format({ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND});
        this->Output("y")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT})
            .Format({ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND});
        this->Attr("reduction").AttrType(OPTIONAL).String("mean");
        this->Attr("ignore_index").AttrType(OPTIONAL).Int(-100);

        this->SetInferShape(ge::InferShape);

        this->AICore()
            .SetTiling(optiling::TilingFunc);
        this->AICore().AddConfig("ascend910b");

    }
};

OP_ADD(NLLLoss);
}
